import os
import time
from typing import Dict, Any
from langchain_core.runnables import RunnablePassthrough, RunnableLambda
from langchain_core.output_parsers import StrOutputParser

# 导入工具类
from tools.AImedal import aiMedal, getMedal


class PromptGenerator:
    """提示词生成器"""
    
    def __init__(self):
        # 性别对应的背景色
        self.color_map = {
            'M': '蓝色',
            'F': '粉色',
            'O': '紫色'  # 其他性别使用紫色
        }
        
        # 健身水平对应的图案
        self.animal_map = {
            'beginner': '麻雀',
            'intermediate': '猫头鹰', 
            'advanced': '老鹰'
        }
    
    def __call__(self, inputs: Dict[str, str]) -> str:
        """生成Logo提示词"""
        gender = inputs['gender']
        fitness_level = inputs['fitness_level']
        
        # 获取对应的颜色和动物
        background_color = self.color_map.get(gender, '蓝色')
        animal = self.animal_map.get(fitness_level, '麻雀')
        
        # 生成简洁的提示词
        prompt = f"一个简洁的健身logo，{background_color}背景，中心图案是{animal}，现代简约风格，适合健身品牌使用"
        print(f"生成的提示词: {prompt}")
        
        return prompt


class AITaskSubmitter:
    """AI任务提交器"""
    
    def __init__(self):
        self.ai_medal = aiMedal
    
    def __call__(self, inputs: Dict[str, Any]) -> Dict[str, Any]:
        """提交AI作画任务"""
        try:
            # 从inputs中提取prompt字符串
            prompt = inputs['prompt']
            
            result = self.ai_medal.submit_task(prompt)
            print(f"[AITaskSubmitter] 提交任务结果: {result}")
            
            return result
        except Exception as e:
            print(f"[AITaskSubmitter] 发生异常: {str(e)}")
            import traceback
            print(f"[AITaskSubmitter] 异常堆栈: {traceback.format_exc()}")
            return {
                'success': False,
                'message': f'提交任务时发生异常: {str(e)}'
            }


class ResultFetcher:
    """结果获取器"""
    
    def __init__(self, max_retries: int = 30):  # 增加重试次数到30次
        self.max_retries = max_retries
        self.get_medal = getMedal
    
    def __call__(self, inputs: Dict[str, Any]) -> Dict[str, Any]:
        """获取AI作画结果，带有重试机制"""
        try:
            # 从inputs中提取task_result
            task_result = inputs['task_result']
            print(f"[ResultFetcher] 提取的task_result: {task_result}")
            
            if not task_result['success']:
                print(f"[ResultFetcher] 任务提交失败，直接返回: {task_result}")
                return task_result
            
            task_id = task_result['task_id']
            print(f"[ResultFetcher] 开始获取任务ID: {task_id} 的结果")
            
            for i in range(self.max_retries):
                print(f"[ResultFetcher] 第 {i+1} 次尝试获取结果")
                result = self.get_medal.get_result(task_id)
                print(f"[ResultFetcher] 获取结果: {result}")

                if result['success']:
                    print(f"[ResultFetcher] 成功获取结果: {result}")
                    return result
                elif result.get('retry', False) and i < self.max_retries - 1:
                    # 任务还在处理中，等待后重试
                    print(f"[ResultFetcher] 任务进行中，等待3秒后重试")
                    time.sleep(3)  # 等待3秒
                    continue
                elif '任务进行中' in result.get('message', '') and i < self.max_retries - 1:
                    # 任务还在处理中，等待后重试
                    print(f"[ResultFetcher] 任务进行中，等待3秒后重试")
                    time.sleep(3)  # 等待3秒
                    continue
                else:
                    # 任务失败或达到最大重试次数
                    print(f"[ResultFetcher] 任务失败或达到最大重试次数: {result}")
                    return result

            print(f"[ResultFetcher] 达到最大重试次数，返回超时错误")
            return {
                'success': False,
                'message': '获取结果超时，请稍后重试'
            }
        except Exception as e:
            print(f"[ResultFetcher] 发生异常: {str(e)}")
            import traceback
            print(f"[ResultFetcher] 异常堆栈: {traceback.format_exc()}")
            return {
                'success': False,
                'message': f'获取结果时发生异常: {str(e)}'
            }


class LogoGenerationAgent:
    """Logo生成Agent，使用管道语法简化流程"""

    def __init__(self):
        # 创建管道组件
        self.prompt_generator = PromptGenerator()
        self.task_submitter = AITaskSubmitter()
        self.result_fetcher = ResultFetcher()
        
        # 构建管道
        self.chain = (
            RunnablePassthrough.assign(prompt=self.prompt_generator)
            | RunnablePassthrough.assign(task_result=self.task_submitter)
            | RunnablePassthrough.assign(result=self.result_fetcher)
            | (lambda x: x['result'])  # 只返回最终结果
        )

    def generate_logo(self, gender: str, fitness_level: str) -> Dict[str, Any]:
        """完整的Logo生成流程 - 使用管道语法"""
        try:
            inputs = {
                'gender': gender,
                'fitness_level': fitness_level
            }
            result = self.chain.invoke(inputs)
            print(f"[LogoGenerationAgent] 管道执行结果: {result}")
            
            return result
        except Exception as e:
            print(f"[LogoGenerationAgent] 发生异常: {str(e)}")
            import traceback
            print(f"[LogoGenerationAgent] 异常堆栈: {traceback.format_exc()}")
            return {
                'success': False,
                'message': f'Logo生成过程中发生异常: {str(e)}'
            }


# 创建全局Agent实例
logo_agent = LogoGenerationAgent()